import torch
import torch.nn as nn
import torch.nn.functional as F
from model.model_utils import _get_padding_mask, _get_visibility_mask
from cadlib.macro import CMD_ARGS_MASK


class CADLoss(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.n_commands = cfg.n_commands
        self.args_dim = cfg.args_dim + 1
        self.weights = cfg.loss_weights

        self.register_buffer("cmd_args_mask", torch.tensor(CMD_ARGS_MASK))

    def forward(self, output):
        # 获取命令和参数的 logits
        command_logits, args_logits = output["command_logits"], output["args_logits"]
        # 获取目标命令和目标参数
        tgt_commands, tgt_args = output["tgt_commands"], output["tgt_args"]

        # 计算 visibility mask 和 padding mask
        visibility_mask = _get_visibility_mask(tgt_commands, seq_dim=-1)
        padding_mask = _get_padding_mask(tgt_commands, seq_dim=-1, extended=True) * visibility_mask.unsqueeze(-1)
        # 确保 padding_mask 的维度与 command_logits 对齐
        padding_mask = padding_mask[:, :command_logits.shape[1]]  # 使 padding_mask 的维度与 command_logits 对应

        # 获取命令的 mask
        mask = self.cmd_args_mask[tgt_commands.long()]

        # 扩展 padding_mask 使其维度与 command_logits 相匹配
        padding_mask_expanded = padding_mask.unsqueeze(-1).expand(-1, -1, command_logits.size(-1))
        
        if padding_mask_expanded.shape != command_logits.shape:
            print(padding_mask_expanded.shape)
            print(command_logits.shape)
            

        # 计算命令的损失
        loss_cmd = F.cross_entropy(
            command_logits[padding_mask_expanded.bool()].reshape(-1, self.n_commands),  # 应用 mask
            tgt_commands[padding_mask.bool()].reshape(-1).long()
        )

        # 计算参数的损失
        loss_args = F.cross_entropy(args_logits[mask.bool()].reshape(-1, self.args_dim),
                                    tgt_args[mask.bool()].reshape(-1).long() + 1)  # shift due to -1 PAD_VAL

        # 加权损失
        loss_cmd = self.weights["loss_cmd_weight"] * loss_cmd
        loss_args = self.weights["loss_args_weight"] * loss_args

        # 返回最终损失
        res = {"loss_cmd": loss_cmd, "loss_args": loss_args}
        return res


